"""
修复后的数据加载和预处理模块
"""
import os
import numpy as np
import scipy.io
import torch
from config import Config

class DataLoader:
    """修复后的数据加载器类"""
    
    def __init__(self, config=None):
        self.config = config or Config()
        self.endmember_matrix = self._load_endmember_matrix()
    
    def _load_endmember_matrix(self):
        """
        加载端元矩阵D
        
        Returns:
            torch.Tensor: 端元矩阵，形状为(128, k)
        """
        try:
            mat_data = scipy.io.loadmat(self.config.ENDMEMBER_PATH)
            D = mat_data["phasecha"]
            
            if D.shape[0] == 64:
                D_reshaped = D.reshape(64, self.config.K)
            else:
                raise ValueError(f"端元矩阵的行数 {D.shape[0]} 与预期的64不匹配")
            
            # 分离实部和虚部
            D_real = np.real(D_reshaped)
            D_imag = np.imag(D_reshaped)
            D_combined = np.concatenate((D_real, D_imag), axis=0)  # (128, k)
            
            return torch.from_numpy(D_combined).float().to(self.config.DEVICE)
            
        except Exception as e:
            raise RuntimeError(f"加载端元矩阵失败: {e}")
    
    def load_batch_data(self, batch_idx, batch_size=None):
        """
        加载一个批次的数据 - 修复后的版本
        
        Args:
            batch_idx (int): 批次索引
            batch_size (int, optional): 批次大小
            
        Returns:
            torch.Tensor: 数据张量，形状为(batch_size, 128, seq_len)
        """
        if batch_size is None:
            batch_size = self.config.BATCH_SIZE
        
        seq_len = self.config.SEQ_LEN
        
        # 创建批次数据容器：(batch_size, 128, seq_len)
        batch_data = torch.zeros(batch_size, 128, seq_len, device=self.config.DEVICE)
        
        # 为每个样本加载一个时间序列
        for sample_idx in range(batch_size):
            # 计算该样本的起始文件索引
            start_file_idx = batch_idx * batch_size * seq_len + sample_idx * seq_len + 1
            
            # 为该样本加载seq_len个连续的时间步
            sample_data = self._load_single_sequence(start_file_idx, seq_len)
            
            if sample_data is not None:
                batch_data[sample_idx] = sample_data.squeeze(0)  # 移除第一个维度
        
        return batch_data
    
    def _load_single_sequence(self, start_idx, seq_len):
        """
        加载单个样本的时间序列数据
        
        Args:
            start_idx (int): 起始文件索引
            seq_len (int): 序列长度
            
        Returns:
            torch.Tensor: 序列数据，形状为(1, 128, seq_len)
        """
        result_matrix = np.zeros((8, 8, seq_len), dtype=np.complex64)
        
        for t in range(seq_len):
            file_idx = start_idx + t
            file_path = os.path.join(
                self.config.DATA_FOLDER,
                f"{self.config.FILE_PREFIX}{file_idx}{self.config.FILE_EXTENSION}"
            )
            
            if os.path.exists(file_path):
                try:
                    mat_data = scipy.io.loadmat(file_path)
                    if "p" in mat_data:
                        Y = mat_data["p"]
                        if Y.size == 64:
                            Y_reshaped = Y.reshape(8, 8)
                            result_matrix[:, :, t] = Y_reshaped
                        else:
                            print(f"警告: {file_path} 中的数据大小不是64")
                    else:
                        print(f"警告: {file_path} 中未找到 'p' 字段")
                except Exception as e:
                    print(f"错误: 读取文件 {file_path} 失败: {e}")
            else:
                print(f"警告: 文件不存在: {file_path}")
                # 如果文件不存在，使用零填充
                result_matrix[:, :, t] = np.zeros((8, 8), dtype=np.complex64)
        
        return self._process_complex_data(result_matrix)
    
    def load_sequence_data(self, start_idx, seq_len=None):
        """
        加载单个序列数据（用于评估）
        
        Args:
            start_idx (int): 起始索引
            seq_len (int, optional): 序列长度
            
        Returns:
            torch.Tensor: 数据张量，形状为(1, 128, seq_len)
        """
        if seq_len is None:
            seq_len = self.config.SEQ_LEN
        
        return self._load_single_sequence(start_idx, seq_len)
    
    def _process_complex_data(self, complex_matrix):
        """
        处理复数数据，分离实部和虚部
        
        Args:
            complex_matrix (np.ndarray): 复数矩阵，形状为(8, 8, time_steps)
            
        Returns:
            torch.Tensor: 处理后的张量，形状为(1, 128, time_steps)
        """
        time_steps = complex_matrix.shape[2]
        
        # 分离实部和虚部
        Y_real = np.real(complex_matrix).reshape(64, time_steps)
        Y_imag = np.imag(complex_matrix).reshape(64, time_steps)
        
        # 拼接实部和虚部
        Y_combined = np.concatenate((Y_real, Y_imag), axis=0)  # (128, time_steps)
        
        # 转换为PyTorch张量并添加批次维度
        Y_tensor = torch.from_numpy(Y_combined).float().to(self.config.DEVICE).unsqueeze(0)
        
        return Y_tensor
    
    def get_endmember_matrix(self):
        """
        获取端元矩阵
        
        Returns:
            torch.Tensor: 端元矩阵
        """
        return self.endmember_matrix
    
    def get_total_sequences(self):
        """
        计算总的可用序列数量
        
        Returns:
            int: 总序列数量
        """
        # 检查数据文件夹中的文件数量
        file_count = 0
        for i in range(1, self.config.TOTAL_DATA + 1):
            file_path = os.path.join(
                self.config.DATA_FOLDER,
                f"{self.config.FILE_PREFIX}{i}{self.config.FILE_EXTENSION}"
            )
            if os.path.exists(file_path):
                file_count += 1
        
        # 计算可以构成多少个完整的序列
        max_sequences = file_count // self.config.SEQ_LEN
        return max_sequences 